# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import math
import logging
from collections import OrderedDict

import google.protobuf.json_format as json_format
import pandas
import statistics
import tabulate

from coin.strategy.marketsim.core import spec_pb2


def aggregate_sim_stat(results, aggregate_key_func=None, show_stat_only=True):
  results = sorted(results, key=lambda r: r.tag)

  dict_results = []
  for res in results:
    res = json_format.MessageToDict(res,
                                    including_default_value_fields=True,
                                    preserving_proto_field_name=True)
    refined = res['product_stats'][0]
    refined['name'] = res['tag']
    refined['return_bp'] = refined['return_per_volume'] * 10000.
    dict_results.append(refined)

  if not show_stat_only:
    try:
      # headers = [f.name for f in spec_pb2.PerProductSimStatProto.DESCRIPTOR.fields]
      headers = [
          'name',
          'gross_pnl',
          'gross_fee',
          'num_orders',
          'num_filled_orders',
          'order_fill_ratio',
          'volume',
          'fill_ratio',
          'taker_ratio',
          'return_bp'
      ]
      rows = [[r[h] for h in headers] for r in dict_results]
      print(tabulate.tabulate(rows, headers=headers))
    except KeyError:
      pass
    except Exception:
      logging.exception('Exception occured')
    print('---')

  print_sim_result_stat(dict_results, aggregate_key_func)


def print_sim_result_stat(results, key_func=None):
  # Assumes results is
  # (1) list of dicts
  # (2) each dict has 'name' key
  # (3) last 8 chracter of 'name' key should be date in YYYYMMDD format.

  if key_func is None:
    key_func = lambda r: (r['name'][:-9], r['name'][-8:])

  times = set()
  agg = {}
  for r in results:
    param_str, time_str = key_func(r)
    times.add(time_str)
    if param_str not in agg:
      agg[param_str] = {}
    if time_str not in agg[param_str]:
      agg[param_str][time_str] = []
    agg[param_str][time_str].append(r)

  params = sorted(agg.keys())
  df_dict = OrderedDict(time=[])
  for p in params:
    df_dict[p] = []

  agg_pnl = {p: [] for p in params}
  agg_fee = {p: [] for p in params}
  agg_volume = {p: [] for p in params}
  agg_cnt = {p: [] for p in params}

  for date_str in sorted(times):
    df_dict['time'].append(date_str)

    for param in params:
      pnl = sum([r['gross_pnl'] for r in agg[param][date_str]])
      fee = sum([r['gross_fee'] for r in agg[param][date_str]])
      volume = sum([r['volume'] for r in agg[param][date_str]])
      cnt = sum([int(r['num_filled_orders']) for r in agg[param][date_str]])

      df_dict[param].append(pnl)
      agg_pnl[param].append(pnl)
      agg_fee[param].append(fee)
      agg_volume[param].append(volume)
      agg_cnt[param].append(cnt)

  df_dict['time'] += [
      '---', 'mean', 'stdev', 'sharpe', 'volume_mean', 'cnt_mean', 'ret_bp', 'fee_bp'
  ]
  for param in params:
    pnl_mean = statistics.mean(agg_pnl[param])
    if len(agg_pnl[param]) > 1:
      pnl_stdev = statistics.stdev(agg_pnl[param])
      pnl_sharpe = pnl_mean / pnl_stdev if pnl_stdev else None
    else:
      pnl_stdev = None
      pnl_sharpe = None

    volume_mean = statistics.mean(agg_volume[param])
    cnt_mean = statistics.mean(agg_cnt[param])
    if sum(agg_volume[param]) != 0:
      ret_bp = sum(agg_pnl[param]) / sum(agg_volume[param]) * 1e4
      fee_bp = sum(agg_fee[param]) / sum(agg_volume[param]) * 1e4
    else:
      ret_bp = None
      fee_bp = None
    df_dict[param] += [None, pnl_mean, pnl_stdev, pnl_sharpe, volume_mean, cnt_mean, ret_bp, fee_bp]

  df = pandas.DataFrame(df_dict)
  df = df.set_index('time')
  pandas.options.display.float_format = lambda x: '' if math.isnan(x) else '{:,.4f}'.format(x)
  with pandas.option_context('display.max_rows',
                             None,
                             'display.max_columns',
                             None,
                             'display.width',
                             250):
    print(df)
